package com.tinyproxy.local;

import com.tinyproxy.common.Cfg;
import com.tinyproxy.common.ExchangeHandler;
import com.tinyproxy.common.Kits;
import com.tinyproxy.local.socks.TinySocks;
import io.netty.bootstrap.Bootstrap;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelOption;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.socksx.SocksMessage;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static com.tinyproxy.common.CodecKind.DECRYPT;
import static com.tinyproxy.common.CodecKind.ENCRYPT;

/**
 * Handler for connecting to remote tiny server
 */
public final class ConnectRemoteHandler extends SimpleChannelInboundHandler<SocksMessage> {

    private final Cfg cfg;

    final Logger log = LoggerFactory.getLogger(getClass());

    private final byte[] staticToken;

    public ConnectRemoteHandler(Cfg cfg) {
        this.cfg = cfg;
        this.staticToken = Kits.toSha256(Kits.asBytes(cfg.getToken()));
    }

    @Override
    public void channelRead0(final ChannelHandlerContext ctx, final SocksMessage message) throws Exception {
        TinySocks socks = TinySocks.instance(message);
        if (socks.invalid()) {
            ctx.close();
            return;
        }

        final Channel inChannel = ctx.channel();
        Bootstrap bootstrap = new Bootstrap();
        bootstrap.group(ctx.channel().eventLoop());
        bootstrap.channel(NioSocketChannel.class);
        bootstrap.option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 10000);
        bootstrap.option(ChannelOption.SO_KEEPALIVE, true);
        bootstrap.handler(new AuthHandler(staticToken, socks, (outCtx, success, token) -> {
            if (success) {
                inChannel.writeAndFlush(socks.respSuccess());

                inChannel.pipeline().addLast(new ExchangeHandler(token, outCtx.channel(), ENCRYPT));
                outCtx.pipeline().addLast(new ExchangeHandler(token, inChannel, DECRYPT));

                inChannel.pipeline().remove(ConnectRemoteHandler.this);
            } else {
                log.info("--Auth failed!!");
                inChannel.writeAndFlush(socks.respReject());
                Kits.closeOnFlush(ctx.channel());
            }
        }));

        bootstrap.connect(cfg.getRemoteHost(), cfg.getRemotePort()).addListener((future) -> {
            if (!future.isSuccess()) {
                // Close the connection if fails occurred.
                log.info("--Close the connection if fails occurred!");
                inChannel.writeAndFlush(socks.respReject());
                Kits.closeOnFlush(ctx.channel());
            }
        });
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        Kits.closeOnFlush(ctx.channel());
    }
}
